import json
import time
import uuid
from typing import Dict, Any, Optional, AsyncIterator, List
from datetime import datetime, timedelta
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI, Request, Response, HTTPException, Header, Depends
from fastapi.responses import StreamingResponse, FileResponse
from loguru import logger
import os
import asyncio
from collections import deque
from time import monotonic

from config import Config

try:
    from fastapi.middleware.cors import CORSMiddleware
except Exception:  # pragma: no cover - optional
    CORSMiddleware = None  # type: ignore

# Prometheus（可选）
try:
    from prometheus_client import (
        Counter,
        Histogram,
        generate_latest,
        CONTENT_TYPE_LATEST,
        CollectorRegistry,
    )  # type: ignore
except Exception:  # pragma: no cover
    Counter = None  # type: ignore
    Histogram = None  # type: ignore
    generate_latest = None  # type: ignore
    CONTENT_TYPE_LATEST = "text/plain"  # type: ignore
    CollectorRegistry = None  # type: ignore

# 共享HTTP客户端（在应用生命周期内复用连接）
_async_client: Optional[httpx.AsyncClient] = None

# hop-by-hop 头部列表
HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}

# 简单限流状态（进程内）
_rate_limiter_windows: Dict[str, deque] = {}

# 指标占位（避免模块被二次导入时重复注册）
REQUEST_COUNTER = None
RESPONSE_COUNTER = None
RATE_LIMIT_COUNTER = None
UPSTREAM_ERRORS_COUNTER = None
LATENCY_HIST = None
_METRICS_REGISTRY = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # startup
    os.makedirs(Config.LOGS_DIR, exist_ok=True)
    os.makedirs(Config.REQUEST_LOGS_DIR, exist_ok=True)
    logger.remove()
    logger.add(
        Config.LOG_FILE,
        rotation=Config.LOG_ROTATION,
        retention=Config.LOG_RETENTION,
        level=Config.LOG_LEVEL,
        encoding=Config.LOG_FILE_ENCODING,
    )
    Config.validate()

    global _async_client
    _async_client = httpx.AsyncClient(timeout=Config.REQUEST_TIMEOUT)

    if Config.CORS_ENABLED and CORSMiddleware is not None:
        origins = [o.strip() for o in (Config.CORS_ALLOW_ORIGINS or "*").split(",")]
        methods = [m.strip() for m in (Config.CORS_ALLOW_METHODS or "*").split(",")]
        headers = [h.strip() for h in (Config.CORS_ALLOW_HEADERS or "*").split(",")]
        app.add_middleware(
            CORSMiddleware,
            allow_origins=origins,
            allow_credentials=Config.CORS_ALLOW_CREDENTIALS,
            allow_methods=methods,
            allow_headers=headers,
        )

    # 初始化指标（仅一次）
    if Config.METRICS_ENABLED and Counter and Histogram and CollectorRegistry:
        global REQUEST_COUNTER, RESPONSE_COUNTER, RATE_LIMIT_COUNTER, UPSTREAM_ERRORS_COUNTER, LATENCY_HIST, _METRICS_REGISTRY
        if REQUEST_COUNTER is None:
            _METRICS_REGISTRY = CollectorRegistry()
            REQUEST_COUNTER = Counter("proxy_requests_total", "总请求数", ["method"], registry=_METRICS_REGISTRY)
            RESPONSE_COUNTER = Counter("proxy_responses_total", "总响应数", ["status_code"], registry=_METRICS_REGISTRY)
            RATE_LIMIT_COUNTER = Counter("proxy_rate_limited_total", "限流次数", registry=_METRICS_REGISTRY)
            UPSTREAM_ERRORS_COUNTER = Counter("proxy_upstream_errors_total", "上游错误次数", registry=_METRICS_REGISTRY)
            LATENCY_HIST = Histogram("proxy_request_latency_seconds", "请求延迟(秒)", registry=_METRICS_REGISTRY)

    try:
        yield
    finally:
        # shutdown
        _client_to_close = _async_client
        _async_client = None
        if _client_to_close is not None:
            await _client_to_close.aclose()

app = FastAPI(title="OpenAI API Proxy", version="1.0.0", lifespan=lifespan)

class RequestLogger:
    """请求日志记录器"""
    
    def __init__(self):
        self.logs_dir = Config.REQUEST_LOGS_DIR
        os.makedirs(self.logs_dir, exist_ok=True)
    
    async def _write_json(self, file_path: str, data: Any):
        def _write():
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
        await asyncio.to_thread(_write)

    # 敏感字段脱敏
    def _sanitize_body(self, value: Any) -> Any:
        try:
            sensitive_keys = {
                "authorization", "api_key", "apikey", "password",
                "access_token", "refresh_token", "token", "secret",
                "client_secret", "session", "cookie", "set-cookie"
            }
            if isinstance(value, dict):
                sanitized = {}
                for k, v in value.items():
                    if isinstance(k, str) and k.lower() in sensitive_keys:
                        sanitized[k] = "***"
                    else:
                        sanitized[k] = self._sanitize_body(v)
                return sanitized
            if isinstance(value, list):
                return [self._sanitize_body(v) for v in value]
            return value
        except Exception:
            return value

    async def _cleanup_logs(self):
        """清理过期或超量的日志文件"""
        try:
            files = [
                os.path.join(self.logs_dir, f)
                for f in os.listdir(self.logs_dir)
                if f.endswith(".json")
            ]
            if not files:
                return
            if Config.REQUEST_LOG_RETENTION_DAYS > 0:
                cutoff = datetime.now() - timedelta(days=Config.REQUEST_LOG_RETENTION_DAYS)
                for fp in files:
                    try:
                        mtime = datetime.fromtimestamp(os.path.getmtime(fp))
                        if mtime < cutoff:
                            os.remove(fp)
                    except Exception:
                        continue
            files = [
                os.path.join(self.logs_dir, f)
                for f in os.listdir(self.logs_dir)
                if f.endswith(".json")
            ]
            if len(files) > Config.REQUEST_LOG_MAX_FILES:
                files.sort(key=lambda p: os.path.getmtime(p))
                to_delete = len(files) - Config.REQUEST_LOG_MAX_FILES
                for fp in files[:to_delete]:
                    try:
                        os.remove(fp)
                    except Exception:
                        continue
        except Exception as e:
            logger.warning(f"日志清理失败: {e}")
    
    def _mask_headers(self, headers: Dict[str, Any]) -> Dict[str, Any]:
        masked = dict(headers)
        for k in list(masked.keys()):
            kl = k.lower()
            if kl in ("authorization", "proxy-authorization", "x-api-key"):
                masked[k] = "***"
        return masked

    def _truncate(self, text: str) -> str:
        limit = max(0, Config.MAX_BODY_LOG_CHARS)
        if limit and len(text) > limit:
            return text[:limit] + "...<truncated>"
        return text
    
    async def log_request(self, request_id: str, method: str, url: str, headers: Dict, body: Any):
        log_data = {
            "request_id": request_id,
            "timestamp": datetime.now().isoformat(),
            "method": method,
            "url": url,
            "headers": self._mask_headers(headers),
        }
        if Config.LOG_REQUEST_BODIES:
            try:
                if isinstance(body, (dict, list)):
                    log_data["body"] = self._sanitize_body(body)
                elif isinstance(body, (bytes, bytearray)):
                    log_data["body"] = self._truncate((body or b"").decode("utf-8", errors="replace"))
                elif isinstance(body, str):
                    # 尝试JSON解析再脱敏
                    try:
                        parsed = json.loads(body)
                        log_data["body"] = self._sanitize_body(parsed)
                    except Exception:
                        log_data["body"] = self._truncate(body)
                else:
                    log_data["body"] = str(body)
            except Exception:
                log_data["body"] = "<unloggable body>"
        
        log_file = os.path.join(self.logs_dir, f"{request_id}_request.json")
        await self._write_json(log_file, log_data)
        asyncio.create_task(self._cleanup_logs())
        logger.info(f"请求已记录: {request_id}")
    
    async def log_response(self, request_id: str, status_code: int, headers: Dict, body: Any):
        if isinstance(body, dict) and body.get("type") == "stream_response":
            if "error" in body:
                log_data = {
                    "request_id": request_id,
                    "timestamp": datetime.now().isoformat(),
                    "status_code": status_code,
                    "response_type": "stream_error",
                    "error": body.get("error", ""),
                    "raw_content_preview": self._truncate(body.get("raw_content", "")),
                    "headers": self._mask_headers(headers)
                }
            else:
                log_data = {
                    "request_id": request_id,
                    "timestamp": datetime.now().isoformat(),
                    "status_code": status_code,
                    "response_type": "stream",
                    "model": body.get("model", ""),
                    "content": self._truncate(body.get("extracted_content", "")),
                    "stop_reason": body.get("stop_reason", ""),
                    "usage": body.get("usage", {}),
                    "raw_events_count": len(body.get("raw_events", [])),
                    "headers": self._mask_headers(headers)
                }
        else:
            log_data = {
                "request_id": request_id,
                "timestamp": datetime.now().isoformat(),
                "status_code": status_code,
                "response_type": "normal",
                "headers": self._mask_headers(headers)
            }
            if Config.LOG_RESPONSE_BODIES:
                if isinstance(body, (bytes, bytearray)):
                    text = (body or b"").decode("utf-8", errors="replace")
                    try:
                        parsed = json.loads(text)
                        log_data["body"] = self._sanitize_body(parsed)
                    except Exception:
                        log_data["body"] = self._truncate(text)
                elif isinstance(body, (dict, list)):
                    log_data["body"] = self._sanitize_body(body)
                else:
                    # str 或其他
                    text = str(body)
                    try:
                        parsed = json.loads(text)
                        log_data["body"] = self._sanitize_body(parsed)
                    except Exception:
                        log_data["body"] = self._truncate(text)
        
        log_file = os.path.join(self.logs_dir, f"{request_id}_response.json")
        await self._write_json(log_file, log_data)
        asyncio.create_task(self._cleanup_logs())
        logger.info(f"响应已记录: {request_id}")

# 计时工具
class _LatencyTimer:
    def __enter__(self):
        self.start = monotonic()
        return self
    def __exit__(self, exc_type, exc, tb):
        self.elapsed = monotonic() - self.start

request_logger = RequestLogger()

async def _require_auth(authorization: Optional[str] = Header(default=None, alias="X-Proxy-Token")):
    if Config.PROXY_AUTH_TOKEN:
        if not authorization or authorization != Config.PROXY_AUTH_TOKEN:
            raise HTTPException(status_code=401, detail="未授权")

@app.middleware("http")
async def log_requests(request: Request, call_next):
    start_time = time.time()
    request_id = str(uuid.uuid4())
    request.state.request_id = request_id

    # 限流：按IP或缺省键
    client_ip = request.client.host if request.client else "unknown"
    bucket_key = client_ip
    now = monotonic()
    window = _rate_limiter_windows.setdefault(bucket_key, deque())
    while window and now - window[0] > Config.RATE_LIMIT_WINDOW_SECONDS:
        window.popleft()
    if len(window) >= Config.RATE_LIMIT_REQUESTS:
        if RATE_LIMIT_COUNTER and Config.METRICS_ENABLED:
            RATE_LIMIT_COUNTER.inc()
        logger.warning(f"限流触发: {client_ip}")
        raise HTTPException(status_code=429, detail="请求过多，请稍后再试")
    window.append(now)

    if REQUEST_COUNTER and Config.METRICS_ENABLED:
        REQUEST_COUNTER.labels(method=request.method.upper()).inc()

    logger.info(f"开始处理请求: {request_id} - {request.method} {request.url}")
    with _LatencyTimer() as t:
        response = await call_next(request)
    # 将请求ID注入响应头，便于排查
    try:
        response.headers["X-Request-ID"] = request_id
    except Exception:
        pass

    if LATENCY_HIST and Config.METRICS_ENABLED:
        LATENCY_HIST.observe(t.elapsed)

    process_time = time.time() - start_time
    logger.info(f"请求完成: {request_id} - 耗时: {process_time:.2f}秒")

    if RESPONSE_COUNTER and Config.METRICS_ENABLED:
        try:
            RESPONSE_COUNTER.labels(status_code=str(response.status_code)).inc()
        except Exception:
            pass

    return response

async def get_request_body(request: Request) -> Any:
    try:
        content_length = request.headers.get("content-length")
        if content_length and int(content_length) > Config.MAX_REQUEST_BODY_BYTES:
            raise HTTPException(status_code=413, detail="请求体过大")
        content_type = request.headers.get("content-type", "").lower()
        if content_type.startswith("application/json"):
            return await request.json()
        return await request.body()
    except HTTPException:
        raise
    except Exception as e:
        logger.warning(f"获取请求体失败: {e}")
        return None

def parse_stream_content(full_content: str) -> Dict[str, Any]:
    try:
        parsed_data = {
            "type": "stream_response",
            "raw_events": [],
            "extracted_content": "",
            "model": "",
            "usage": {},
            "stop_reason": ""
        }
        lines = full_content.strip().split('\n')
        for line in lines:
            line = line.strip()
            if not line:
                continue
            if line.startswith('data:'):
                data_content = line[5:].strip()
                if data_content == '[DONE]':
                    break
                try:
                    event_data = json.loads(data_content)
                    parsed_data["raw_events"].append(event_data)
                    if event_data.get('type') == 'message_start':
                        model = event_data.get('message', {}).get('model', '')
                        parsed_data["model"] = model
                    elif event_data.get('type') == 'content_block_delta':
                        delta = event_data.get('delta', {})
                        if delta.get('type') == 'text_delta':
                            text = delta.get('text', '')
                            parsed_data["extracted_content"] += text
                    elif event_data.get('type') == 'message_delta':
                        stop_reason = event_data.get('delta', {}).get('stop_reason', '')
                        usage = event_data.get('usage', {})
                        parsed_data["stop_reason"] = stop_reason
                        parsed_data["usage"] = usage
                    elif 'choices' in event_data:
                        choices = event_data.get('choices', [])
                        for choice in choices:
                            if 'delta' in choice:
                                delta = choice['delta']
                                if 'content' in delta:
                                    parsed_data["extracted_content"] += delta['content']
                                if 'model' in delta:
                                    parsed_data["model"] = delta['model']
                            elif 'message' in choice:
                                message = choice['message']
                                if 'content' in message:
                                    parsed_data["extracted_content"] += message['content']
                        if 'usage' in event_data:
                            parsed_data["usage"] = event_data['usage']
                    elif 'output' in event_data:
                        output = event_data.get('output', {})
                        if 'text' in output:
                            parsed_data["extracted_content"] += output['text']
                        if 'model' in output:
                            parsed_data["model"] = output['model']
                        if 'usage' in output:
                            parsed_data["usage"] = output['usage']
                        if 'finish_reason' in output:
                            parsed_data["stop_reason"] = output['finish_reason']
                except json.JSONDecodeError:
                    continue
        logger.info(f"解析完成: 模型={parsed_data['model']}, 内容长度={len(parsed_data['extracted_content'])}, 事件数量={len(parsed_data['raw_events'])}")
        return parsed_data
    except Exception as e:
        logger.error(f"解析流式响应失败: {e}")
        preview = full_content[: Config.MAX_STREAM_LOG_BYTES]
        return {
            "type": "stream_response",
            "error": f"解析失败: {str(e)}",
            "raw_content": preview + ("..." if len(full_content) > len(preview) else "")
        }


def _clean_headers(headers: Dict[str, str]) -> Dict[str, str]:
    cleaned = dict(headers)
    cleaned.pop("host", None)
    cleaned.pop("content-length", None)
    for h in list(cleaned.keys()):
        if h.lower() in HOP_BY_HOP_HEADERS:
            cleaned.pop(h, None)
    return cleaned

async def _build_outgoing_headers(incoming_headers: Dict[str, str]) -> Dict[str, str]:
    headers = _clean_headers(incoming_headers)
    if Config.OVERRIDE_AUTHORIZATION and Config.OPENAI_API_KEY:
        headers["Authorization"] = f"Bearer {Config.OPENAI_API_KEY}"
    return headers

async def _tee_stream_and_log(response: httpx.Response, request_id: str) -> AsyncIterator[bytes]:
    chunks: List[bytes] = []
    total = 0
    async for chunk in response.aiter_bytes():
        if Config.LOG_STREAM_CONTENT and total < Config.MAX_STREAM_LOG_BYTES:
            chunks.append(chunk)
            total += len(chunk)
        yield chunk
    if Config.LOG_STREAM_CONTENT and chunks:
        async def _parse_and_log():
            try:
                full_content = b"".join(chunks).decode("utf-8", errors="replace")
                parsed_content = parse_stream_content(full_content)
                await request_logger.log_response(request_id, response.status_code, dict(response.headers), parsed_content)
            except Exception as e:
                logger.warning(f"流式日志记录失败: {e}")
        asyncio.create_task(_parse_and_log())

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD"])
async def proxy_openai(request: Request, path: str, _: None = Depends(_require_auth)):
    request_id = getattr(request.state, "request_id", str(uuid.uuid4()))
    target_url = f"{Config.OPENAI_BASE_URL}/{path}"
    headers = await _build_outgoing_headers(dict(request.headers))
    if not Config.OVERRIDE_AUTHORIZATION and "authorization" not in {k.lower() for k in headers.keys()} and Config.OPENAI_API_KEY:
        headers["Authorization"] = f"Bearer {Config.OPENAI_API_KEY}"
    body = await get_request_body(request)
    asyncio.create_task(request_logger.log_request(request_id, request.method, target_url, headers, body))
    params = dict(request.query_params)
    try:
        assert _async_client is not None
        client = _async_client
        resp = await client.request(
            method=request.method,
            url=target_url,
            headers=headers,
            params=params,
            content=body if isinstance(body, (bytes, bytearray)) else None,
            json=body if isinstance(body, dict) else None,
        )
        resp_headers = _clean_headers(dict(resp.headers))
        if resp.headers.get("content-type", "").startswith("text/event-stream"):
            return StreamingResponse(
                _tee_stream_and_log(resp, request_id),
                status_code=resp.status_code,
                headers=resp_headers,
                media_type=resp.headers.get("content-type")
            )
        else:
            content = resp.content
            asyncio.create_task(request_logger.log_response(request_id, resp.status_code, dict(resp.headers), content))
            return Response(
                content=content,
                status_code=resp.status_code,
                headers=resp_headers,
                media_type=resp.headers.get("content-type")
            )
    except httpx.RequestError as e:
        if UPSTREAM_ERRORS_COUNTER and Config.METRICS_ENABLED:
            UPSTREAM_ERRORS_COUNTER.inc()
        logger.error(f"请求失败: {request_id} - {e}")
        raise HTTPException(status_code=500, detail=f"代理请求失败: {str(e)}")
    except Exception as e:
        logger.error(f"未知错误: {request_id} - {e}")
        raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")

# 健康检查：可选深度检查上游连接
@app.get("/health")
async def health_check(deep: Optional[bool] = False):
    result = {"status": "healthy", "timestamp": datetime.now().isoformat()}
    if deep:
        try:
            assert _async_client is not None
            headers = {}
            if Config.OPENAI_API_KEY:
                headers["Authorization"] = f"Bearer {Config.OPENAI_API_KEY}"
            url = f"{Config.OPENAI_BASE_URL.rstrip('/')}/v1/models"
            resp = await _async_client.get(url, headers=headers, timeout=5.0)
            result["upstream_status"] = resp.status_code
            result["upstream_ok"] = 200 <= resp.status_code < 500  # 4xx也视为可达
        except Exception as e:
            result["upstream_ok"] = False
            result["error"] = str(e)
            result["status"] = "degraded"
    return result

# 日志列表端点
@app.get("/logs")
async def get_logs(request: Request, _: None = Depends(_require_auth)):
    logs_dir = Config.REQUEST_LOGS_DIR
    if not os.path.exists(logs_dir):
        return {"total": 0, "page": 1, "page_size": 0, "items": []}

    # 查询参数
    q = request.query_params
    page = max(1, int(q.get("page", 1)))
    page_size = max(1, min(100, int(q.get("page_size", 20))))
    type_filter = q.get("type")  # request / response / None
    rid_filter = q.get("request_id")  # 前缀匹配
    since_s = q.get("since")
    until_s = q.get("until")

    def parse_time(s: Optional[str]) -> Optional[datetime]:
        if not s:
            return None
        try:
            return datetime.fromisoformat(s)
        except Exception:
            return None

    since = parse_time(since_s)
    until = parse_time(until_s)

    items: List[Dict[str, Any]] = []
    for file in os.listdir(logs_dir):
        if not file.endswith(".json"):
            continue
        if type_filter == "request" and not file.endswith("_request.json"):
            continue
        if type_filter == "response" and not file.endswith("_response.json"):
            continue
        if rid_filter and not file.startswith(rid_filter):
            continue
        file_path = os.path.join(logs_dir, file)
        stat = os.stat(file_path)
        mtime = datetime.fromtimestamp(stat.st_mtime)
        if since and mtime < since:
            continue
        if until and mtime > until:
            continue
        items.append({
            "filename": file,
            "size": stat.st_size,
            "modified": mtime.isoformat(),
            "type": "request" if file.endswith("_request.json") else ("response" if file.endswith("_response.json") else "unknown"),
            "request_id": file.replace("_request.json", "").replace("_response.json", "")
        })

    # 排序与分页
    items.sort(key=lambda x: x["modified"], reverse=True)
    total = len(items)
    start = (page - 1) * page_size
    end = start + page_size
    page_items = items[start:end]

    return {"total": total, "page": page, "page_size": page_size, "items": page_items}

def _safe_join_logs_dir(filename: str) -> Optional[str]:
    if not filename.endswith(".json"):
        return None
    candidate = os.path.abspath(os.path.join(Config.REQUEST_LOGS_DIR, filename))
    base = os.path.abspath(Config.REQUEST_LOGS_DIR)
    if os.path.commonprefix([candidate, base]) != base:
        return None
    if not os.path.exists(candidate):
        return None
    return candidate

@app.get("/logs/{request_id}")
async def get_log_pair(request_id: str, _: None = Depends(_require_auth)):
    req_file = f"{request_id}_request.json"
    resp_file = f"{request_id}_response.json"
    req_path = _safe_join_logs_dir(req_file)
    resp_path = _safe_join_logs_dir(resp_file)
    result: Dict[str, Any] = {"request_id": request_id}
    if req_path:
        with open(req_path, "r", encoding="utf-8") as f:
            result["request"] = json.load(f)
    if resp_path:
        with open(resp_path, "r", encoding="utf-8") as f:
            result["response"] = json.load(f)
    if "request" not in result and "response" not in result:
        raise HTTPException(status_code=404, detail="未找到该请求ID的日志")
    return result

@app.get("/logs/file/{filename}")
async def download_log_file(filename: str, _: None = Depends(_require_auth)):
    path = _safe_join_logs_dir(filename)
    if not path:
        raise HTTPException(status_code=404, detail="文件不存在或不允许访问")
    return FileResponse(path, media_type="application/json", filename=filename)

# 指标端点
@app.get("/metrics")
async def metrics(_: None = Depends(_require_auth) if Config.METRICS_AUTH_REQUIRED else None):
    if not Config.METRICS_ENABLED or generate_latest is None or _METRICS_REGISTRY is None:
        raise HTTPException(status_code=404, detail="未启用指标")
    data = generate_latest(_METRICS_REGISTRY)  # type: ignore
    return Response(content=data, media_type=CONTENT_TYPE_LATEST)

if __name__ == "__main__":
    import uvicorn
    os.makedirs(Config.LOGS_DIR, exist_ok=True)
    os.makedirs(Config.REQUEST_LOGS_DIR, exist_ok=True)
    Config.print_config()
    logger.info("启动OpenAI API代理服务器...")
    logger.info(f"OpenAI Base URL: {Config.OPENAI_BASE_URL}")
    uvicorn.run(
        "main:app",
        host=Config.HOST,
        port=Config.PORT,
        reload=Config.DEBUG,
        log_level=Config.LOG_LEVEL.lower()
    ) 